

#include <assert.h>
#include "oneapi/dnnl/dnnl.h"
#include "opdesc.hpp"
#include "primitive_desc_iface.hpp"

#include "c_types_map.hpp"
#include "memory_desc_wrapper.hpp"
#include "utils.hpp"

using namespace std;
using namespace dnnl::impl;
using namespace dnnl::impl::utils;
using namespace dnnl::impl::status;
using namespace dnnl::impl::prop_kind;
using namespace dnnl::impl::alg_kind;
using namespace dnnl::impl::types;

#define VCHECK_SPATIALTF(cond, msg, ...) \
VCONDCHECK(primitive, create, check, eltwise, (cond), \
status::invalid_arguments, msg, ##__VA_ARGS__);

#define VCHECK_SPATIALTF(cond, msg, ...) \
VCONDCHECK(primitive, create, check, eltwise, (cond), \
status::unimplemented, msg, ##__VA_ARGS__);

namespace dnnl {
namespace impl {
status_t spatialtf_desc_init(spatialtf_desc_t *spatialtf_desc, prop_kind_t prop_kind,
         const memory_desc_t *src0_desc,
        const memory_desc_t *src1_desc, const memory_desc_t *dst_desc) {
    const bool is_fwd = one_of(prop_kind, forward_training, forward_inference);
    VCHECK_SPATIALTF(!any_null(spatialtf_desc, dst_desc), VERBOSE_NULL_ARG);

    auto sd = spatialtf_desc_t();
    sd.primitive_kind = primitive_kind::spatialtf;
    sd.prop_kind = prop_kind;
    sd.src1_desc = *src1_desc;
    sd.src0_desc = *src0_desc;
    sd.dst_desc = *dst_desc;
    *spatialtf_desc = sd;

    return status::success;
}
} // namespace impl
} // namespace dnnl

status_t dnnl_spatialtf_grid_generator_forward_primitive_desc_create(primitive_desc_iface_t **primitive_desc_iface, dnnl_engine_t engine,
        dnnl_prop_kind_t prop_kind,
        const_dnnl_memory_desc_t src0_desc,
        const_dnnl_memory_desc_t src1_desc,
        const_dnnl_memory_desc_t dst_desc,
        const primitive_attr_t *attr) {
    auto spatialtf_desc = spatialtf_desc_t();
    CHECK(spatialtf_desc_init(&spatialtf_desc, prop_kind, src0_desc, src1_desc, dst_desc));

    return primitive_desc_create(primitive_desc_iface, engine,(const op_desc_t *)(&spatialtf_desc), nullptr, attr);
}